package branch_and_bound;

import gurobi.GRB;
import gurobi.GRBConstr;
import gurobi.GRBEnv;
import gurobi.GRBException;
import gurobi.GRBLinExpr;
import gurobi.GRBModel;
import gurobi.GRBVar;
import gurobi.GRB.CharAttr;
import gurobi.GRB.DoubleAttr;
import gurobi.GRB.IntAttr;
import gurobi.GRB.StringAttr;

public class Solver {

    private GRBEnv env;
    private GRBModel model;

    private double optimal;
    private double[] variables;

    Solver(GRBEnv env, GRBModel model) {
        this.env = env;
        this.model = model;

        this.optimal = -Double.MAX_VALUE;
        this.variables = new double[this.model.getVars().length];
    }

    public void solve() {
        try {
            for (GRBVar i : model.getVars()) {
                i.set(GRB.CharAttr.VType, GRB.CONTINUOUS);
            }
            solveModel(model);
        } catch (GRBException e) {
            System.out.println("Error code: " + e.getErrorCode() + ". " + e.getMessage());
        }

        log("===========================================");
        log("最优值为：" + optimal);
        log("最优解为：");
        for (int i = 0; i < variables.length; i++) {
            log("x" + i + " = " + variables[i]);
        }
    }

    // 判断变量是否均为整数
    private boolean intVariables(GRBVar[] vars) {
        for (GRBVar i : vars) {
            try {
                double v = i.get(DoubleAttr.X);
                if (v != (int) v) {
                    return false;
                }
            } catch (GRBException e) {
                System.out.println("Error code: " + e.getErrorCode() + ". " + e.getMessage());
            }
        }
        return true;
    }

    // 找到最大的非整数变量
    private int maxNonintVariable(GRBModel model) {
        int index = -1;
        double value = -Double.MAX_VALUE;
        try {
            for (int i = 0; i < model.get(IntAttr.NumVars); i++) {
                double v = model.getVars()[i].get(DoubleAttr.X);
                if (v > value && v != (int) v) {
                    value = v;
                    index = i;
                }
            }
        } catch (GRBException e) {
            System.out.println("Error code: " + e.getErrorCode() + ". " + e.getMessage());
        }
        return index;
    }

    // 复制当前模型
    private GRBModel copyModel(GRBModel model) {
        GRBModel m = null;
        try {
            m = new GRBModel(env);
            m.set("OutputFlag", "0");

            // 复制变量
            m.addVars(model.get(IntAttr.NumVars), GRB.CONTINUOUS);

            // 复制目标函数
            m.setObjective(model.getObjective(), model.get(IntAttr.ModelSense));

            // 复制约束
            for (GRBConstr i : model.getConstrs()) {
                m.addConstr(
                        model.getRow(i),
                        i.get(CharAttr.Sense),
                        i.get(DoubleAttr.RHS),
                        i.get(StringAttr.ConstrName));
            }

        } catch (GRBException e) {
            System.out.println("Error code: " + e.getErrorCode() + ". " + e.getMessage());
        }
        return m;
    }

    // 输出日志
    private void log(String m) {
        System.out.println(">>> " + m);
    }

    private void solveModel(GRBModel m) {
        log("===================");
        try {

            m.optimize();

            log("求解结果: ");
            if (m.get(IntAttr.Status) == GRB.Status.OPTIMAL) {
                double currentSolution = m.get(GRB.DoubleAttr.ObjVal);
                System.out.println(currentSolution);
                for (int i = 0; i < m.getVars().length; i++) {
                    System.out.println(m.getVars()[i].get(DoubleAttr.X));
                }

                if (currentSolution >= optimal) {
                    if (intVariables(m.getVars())) {
                        optimal = currentSolution;
                        for (int i = 0; i < m.getVars().length; i++) {
                            variables[i] = m.getVars()[i].get(DoubleAttr.X);
                        }
                        log("发现一组可行解");
                        log("===================");
                    } else {
                        int index = maxNonintVariable(m);
                        log("开始分支, 分支变量为: " + index);
                        if (index != -1) {
                            generateSubproblems(index, m);
                        } else {
                            log("!!!warning!!!");
                        }
                    }
                } else {
                    log("比当前最优值小，剪去该分支");
                    log("===================");
                }
            } else {
                // System.out.println(m.get(IntAttr.Status));
                log("该分支无解，剪去");
            }
        } catch (GRBException e) {
            System.out.println("Error code: " + e.getErrorCode() + ". " + e.getMessage());
        }
    }

    private void generateSubproblems(int index, GRBModel m) {
        try {
            GRBModel m1 = copyModel(m);
            GRBModel m2 = copyModel(m);

            GRBVar branchVar = m.getVars()[index];
            int v = (int) branchVar.get(DoubleAttr.X);

            GRBLinExpr expr = new GRBLinExpr();
            expr.addTerm(1, branchVar);
            m1.addConstr(expr, GRB.LESS_EQUAL, v, "cut");
            m2.addConstr(expr, GRB.GREATER_EQUAL, v + 1, "cut");

            log("添加约束: x" + index + " <= " + v);
            solveModel(m1);
            log("添加约束: x" + index + " >= " + (v + 1));
            solveModel(m2);
        } catch (GRBException e) {
            System.out.println("Error code: " + e.getErrorCode() + ". " + e.getMessage());
        }
    }
}
